from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.chat_models import ChatTongyi
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
import os

# 设置通义千问API密钥 (从环境变量读取)
# 1、读取环境变量配置的api_key和url
AI_QW_KEY = os.getenv("AI_QW_KEY")
AI_QW_URL =os.getenv("AI_QW_URL")

# 设置通义千问的 API Key
os.environ["DASHSCOPE_API_KEY"] = AI_QW_KEY # 替换为你的 API Key
# 可选：设置通义千问的 API 端点（如果需要）
os.environ["DASHSCOPE_ENDPOINT"] = AI_QW_URL


def pdf_rag(file_path: str, query: str):
    """
    实现PDF文件的RAG问答流程
    :param file_path: PDF文件路径
    :param query: 用户问题
    """
    # 1. 加载PDF文档
    loader = PyPDFLoader(file_path)
    documents = loader.load()
    print(f"已加载 {len(documents)} 页文档")

    # 2. 分割文本
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,  # 每个块500字符
        chunk_overlap=50,  # 块间重叠50字符
        length_function=len
    )
    texts = text_splitter.split_documents(documents)
    print(f"分割为 {len(texts)} 个文本块")

    # 3. 创建向量数据库
    embeddings = DashScopeEmbeddings(
        model="text-embedding-v1",  # 通义千问官方嵌入模型
    )

    # 将文本向量化并存储到ChromaDB
    vector_db = Chroma.from_documents(
        documents=texts,
        embedding=embeddings,
        persist_directory="./chroma_db"  # 向量数据库存储路径
    )
    print("向量数据库创建完成")

    # 4. 创建检索式问答链
    llm = ChatTongyi(model="qwen-turbo")  # 使用通义千问Turbo模型

    # 自定义提示模板
    prompt_template = """
    使用以下上下文片段来回答问题。如果你不知道答案，请如实说明。
    上下文：
    {context}

    问题：{question}
    答案：
    """
    PROMPT = PromptTemplate(
        template=prompt_template,
        input_variables=["context", "question"]
    )

    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=vector_db.as_retriever(search_kwargs={"k": 3}),  # 检索前3个相关片段
        chain_type_kwargs={"prompt": PROMPT},
        return_source_documents=True
    )

    # 5. 执行查询
    result = qa_chain.invoke({"query": query})

    # 6. 打印结果
    print("\n" + "=" * 50)
    print(f"问题: {query}")
    print(f"答案: {result['result']}")
    print("\n来源文档片段:")
    for doc in result['source_documents']:
        print(f"- 页码 {doc.metadata['page']}: {doc.page_content[:100]}...")


if __name__ == "__main__":
    # 使用示例
    pdf_file = "quanliucun.pdf"  # 替换为你的PDF文件路径
    question = "王孬"

    pdf_rag(pdf_file, question)